﻿// CShellcodeDlg.cpp: 实现文件
//

#include "pch.h"
#include "Injector.h"
#include "CShellcodeDlg.h"
#include "afxdialogex.h"
#include "CtrlCode.h"


// CShellcodeDlg 对话框

IMPLEMENT_DYNAMIC(CShellcodeDlg, CDialogEx)

CShellcodeDlg::CShellcodeDlg(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_SHELLCODE_DIALOG, pParent)
    , m_strShellCode(_T(""))
{

}

CShellcodeDlg::~CShellcodeDlg()
{
}

void CShellcodeDlg::DoDataExchange(CDataExchange* pDX)
{
    CDialogEx::DoDataExchange(pDX);
    DDX_Text(pDX, IDC_EDIT_SHELLCODE, m_strShellCode);
    DDX_Control(pDX, IDC_EDIT_DISASM, m_EditDisasm);
    DDX_Control(pDX, IDC_EDIT_SHELLCODE, m_EditInputCode);
}


BEGIN_MESSAGE_MAP(CShellcodeDlg, CDialogEx)
    ON_WM_CLOSE()
    ON_BN_CLICKED(IDC_BTN_CONFIRM, &CShellcodeDlg::OnBnClickedBtnConfirm)
    ON_EN_CHANGE(IDC_EDIT_SHELLCODE, &CShellcodeDlg::OnEnChangeEditShellcode)
    ON_MESSAGE(WM_GETINJECTADDRESS, &CShellcodeDlg::OnRecvCodeInjectAddress)
END_MESSAGE_MAP()


// CShellcodeDlg 消息处理程序

void WriteAsciiToMemory(char* src, char* dst, int len)
{
    char Buff[3] = { 0 };
    char* s = nullptr;
    long l;
    for (int i = 0; i < len; i += 2)
    {
        Buff[0] = src[i];
        Buff[1] = src[i + 1];
        l = strtol(Buff, &s, 16);
        memcpy(dst, &l, 1);
        dst++;
    }
}

BOOL CShellcodeDlg::PreTranslateMessage(MSG* pMsg)
{
    if (pMsg->message == WM_KEYDOWN)
    {
        switch (pMsg->wParam)
        {
        case VK_ESCAPE:
            return TRUE;
        case VK_RETURN:
            return TRUE;
        default:
            break;
        }
    }

    return CDialogEx::PreTranslateMessage(pMsg);
}


void CShellcodeDlg::OnClose()
{
    ::SendMessage(this->GetParent()->m_hWnd, WM_SHELLCODEINJECT, 0, 0);
    CDialogEx::OnClose();
}


void CShellcodeDlg::DisasmBuff(ZydisDecoder* decoder, ZyanU8* data, ZyanUSize length)
{
    ZydisFormatter formatter;
    ZyanU64 RtlAddress = 0x0;
    ZydisDecodedInstruction instruction;
    char buffer[256];
    CString str;

    ZydisFormatterInit(&formatter, ZYDIS_FORMATTER_STYLE_INTEL);
    ZydisFormatterSetProperty(&formatter, ZYDIS_FORMATTER_PROP_FORCE_SEGMENT, ZYAN_TRUE);
    ZydisFormatterSetProperty(&formatter, ZYDIS_FORMATTER_PROP_FORCE_SIZE, ZYAN_TRUE);

    while (ZYAN_SUCCESS(ZydisDecoderDecodeBuffer(decoder, data, length, &instruction)))
    {
        RtlZeroMemory(buffer, 256);
        ZydisFormatterFormatInstruction(&formatter, &instruction, &buffer[0], sizeof(buffer), RtlAddress);
        str += buffer;
        str += "\r\n";
        data += instruction.length;
        length -= instruction.length;
        RtlAddress += instruction.length;
    }
    m_EditDisasm.SetWindowTextW(str);
}

void CShellcodeDlg::OnBnClickedBtnConfirm()
{
    UpdateData(TRUE);

    PSHELLCODE_CONFIG pConfig = (PSHELLCODE_CONFIG)new char[sizeof(SHELLCODE_CONFIG)];
    BSTR pWch;
    int nLen;
    CString strCode;
    PCHAR pCode;
    PCHAR pTemp;

    if (pConfig && m_strShellCode.GetLength())
    {
        RtlZeroMemory(pConfig, sizeof(SHELLCODE_CONFIG));
        pConfig->bConfirm = TRUE;

        //
        //  删除无关字符
        //

        m_strShellCode.Remove('\\');
        m_strShellCode.Remove('0x');
        m_strShellCode.Remove('"');
        m_strShellCode.Remove('x');
        m_strShellCode.Remove(' ');
        m_strShellCode.Remove(',');

        //
        //  转换成多字节
        //

        strCode = m_strShellCode;
        pWch = strCode.GetBuffer(0);
        nLen = strCode.GetLength();
        pCode = new char[nLen];
        pTemp = new char[nLen];
        if (!pCode || !pTemp) goto _None;
        WideCharToMultiByte(CP_ACP, 0, pWch, -1, pTemp, nLen, NULL, FALSE);

        //
        //  将ASCII码直接写入内存
        //
        //  e.g
        //
        //  Memory                          Ascii           Memory          Ascii
        //  45 39 32 32 37 41 36 42 30 30	E9227A6B00  ->  e9 22 7a 6b 00  ?"zk.
        //

        WriteAsciiToMemory(pTemp, pCode, nLen);

        //
        //  发送消息
        //

        pConfig->CodeLen = nLen / 2;
        pConfig->Code = pCode;
        ::SendMessage(this->GetParent()->m_hWnd, WM_SHELLCODEINJECT, (WPARAM)pConfig, 0);
        delete[] pTemp;
        delete[] pCode;
    }
    else
    {
_None:
        ::SendMessage(this->GetParent()->m_hWnd, WM_SHELLCODEINJECT, 0, 0);
    }

    if (pConfig)
        delete[] pConfig;

    EndDialog(0);
}


BOOL CShellcodeDlg::OnInitDialog()
{
    CDialogEx::OnInitDialog();

    m_lpCodeAddress = 0;

    ::SendMessage(this->GetParent()->m_hWnd, WM_GETINJECTADDRESS, (WPARAM)this->m_hWnd, 0);

    CString strCodeAddress;
    strCodeAddress.Format(_T("Address is 0x%X"), m_lpCodeAddress);
    AfxMessageBox(strCodeAddress);

    CFont Font;
    Font.CreatePointFont(150, _T("新宋体"));
    m_EditDisasm.SetFont(&Font);

    m_EditInputCode.SetWindowTextW(_T(""));

    return TRUE;
}


void CShellcodeDlg::OnEnChangeEditShellcode()
{
    UpdateData(TRUE);
    m_EditDisasm.Clear();
    if (!m_strShellCode.GetLength())
    {
        m_EditDisasm.SetWindowTextW(_T(""));
        return;
    }

    CString strCode;
    PWCHAR pWch;
    int nLen;
    PCHAR pCode = nullptr;
    PCHAR pTemp = nullptr;
    ZydisDecoder decoder;

    do
    {
        //
        //  删除无关字符
        //

        m_strShellCode.Remove('\\');
        m_strShellCode.Remove('0x');
        m_strShellCode.Remove('"');
        m_strShellCode.Remove('x');
        m_strShellCode.Remove(' ');
        m_strShellCode.Remove(',');

        //
        //  转换成多字节
        //

        strCode = m_strShellCode;
        pWch = strCode.GetBuffer(0);
        nLen = strCode.GetLength();
        pCode = new char[nLen];
        pTemp = new char[nLen];
        if (!pCode || !pTemp) break;
        WideCharToMultiByte(CP_ACP, 0, pWch, -1, pTemp, nLen, NULL, FALSE);

        //
        //  将ASCII码直接写入内存
        //

        WriteAsciiToMemory(pTemp, pCode, nLen);

#ifdef _WIN64
        ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64);
#else
        ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_COMPAT_32, ZYDIS_ADDRESS_WIDTH_32);
#endif

        DisasmBuff(&decoder, (ZyanU8*)pCode, nLen / 2);

    } while (FALSE);

    if (pTemp)
        delete[] pTemp;

    if (pCode)
        delete[] pCode;
}

LRESULT CShellcodeDlg::OnRecvCodeInjectAddress(WPARAM wParam, LPARAM lParam)
{
    m_lpCodeAddress = (PVOID)wParam;
    return 0;
}
